package org.acghub.langchain4j.message;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Function;

public abstract class AbstractMessageMemory<K, V> implements MessageMemory<K, V> {

    protected final Map<K, List<V>> memory;

    protected final Function<K, List<V>> valueListSupplier;

    public AbstractMessageMemory(Map<K, List<V>> memory, Function<K, List<V>> valueListSupplier) {
        this.memory = memory;
        this.valueListSupplier = valueListSupplier;
    }

    public AbstractMessageMemory() {
        this(new ConcurrentHashMap<>(64), k -> new CopyOnWriteArrayList<>());
    }

    protected List<V> getValueList(K k) {
        return memory.computeIfAbsent(k, valueListSupplier);
    }

    @Override
    public V getLast(K k) {
        List<V> list = getValueList(k);
        if (!list.isEmpty()) {
            return list.get(list.size() - 1);
        }
        return null;
    }

    @Override
    public V getFirst(K k) {
        List<V> list = getValueList(k);
        if (!list.isEmpty()) {
            return list.get(0);
        }
        return null;
    }

    @Override
    public List<V> getHistory(K k) {
        return new ArrayList<>(getValueList(k));
    }

    @Override
    public V get(K k, int index) {
        List<V> list = getValueList(k);
        if (index < list.size()) {
            return list.get(index);
        }
        return null;
    }

    @Override
    public void addLast(K k, V v) {
        getValueList(k).add(v);
    }

    @Override
    public void addFirst(K k, V v) {
        getValueList(k).add(0, v);
    }

    @Override
    public void add(K k, V v, int index) {
        getValueList(k).add(index, v);
    }

    @Override
    public void remove(K k, int index) {
        List<V> list = getValueList(k);
        if (index < list.size()) {
            list.remove(index);
        }
    }

    @Override
    public void removeFirst(K k) {
        List<V> list = getValueList(k);
        if (!list.isEmpty()) {
            list.remove(0);
        }
    }

    @Override
    public void removeLast(K k) {
        List<V> list = getValueList(k);
        if (!list.isEmpty()) {
            list.remove(list.size() - 1);
        }

    }

    @Override
    public void delete(K k) {
        memory.remove(k);
    }
}
